#define _CRT_SECURE_NO_WARNINGS
#define HAVE_REMOTE
#define WIN32

#include <stdio.h>
#include <stdlib.h>
#include <pcap.h>
#include <Packet32.h>
#include <ntddndis.h>
#include <string.h>
#include <time.h>
#include <winsock2.h> // For WSAStartup, WSACleanup, inet_ntoa

#pragma comment(lib, "Packet")
#pragma comment(lib, "wpcap")
#pragma comment(lib, "WS2_32")

// --- Constants ---
#define LOG_FILENAME "network_log.csv"
#define STATS_INTERVAL_SECONDS 60
#define MAC_ADDR_LEN 6
#define IP_ADDR_STR_LEN 16 // Max length for xxx.xxx.xxx.xxx
#define MAC_ADDR_STR_LEN 18 // Max length for XX-XX-XX-XX-XX-XX
#define TIMESTAMP_STR_LEN 20 // Max length for YYYY-MM-DD HH:MM:SS

// --- Data Structures ---

// Ethernet Header (14 bytes)
typedef struct mac_header {
    u_char dest_addr[MAC_ADDR_LEN];
    u_char src_addr[MAC_ADDR_LEN];
    u_short type; // Network byte order
} mac_header;

// IP Header (simplified, minimum 20 bytes, assumes no options for offset calculation)
typedef struct ip_header {
    u_char ver_ihl;       // Version (4 bits) + Internet header length (4 bits)
    u_char tos;           // Type of service
    u_short tlen;         // Total length (network byte order)
    u_short identification; // Identification (network byte order)
    u_short flags_fo;     // Flags (3 bits) + Fragment offset (13 bits) (network byte order)
    u_char ttl;           // Time to live
    u_char proto;         // Protocol
    u_short crc;          // Header checksum (network byte order)
    struct in_addr saddr; // Source address (use struct in_addr for easier handling)
    struct in_addr daddr; // Destination address
    // Options and padding are variable and follow here...
} ip_header;

// Structure to hold statistics for one address (MAC or IP)
typedef struct {
    char address[MAC_ADDR_STR_LEN > IP_ADDR_STR_LEN ? MAC_ADDR_STR_LEN : IP_ADDR_STR_LEN]; // Sufficient space for MAC or IP string
    unsigned long long sent_bytes;
    unsigned long long recv_bytes;
} AddressStat;

// Structure for managing a dynamic list of AddressStat
typedef struct {
    AddressStat* stats;
    size_t count;
    size_t capacity;
} StatsList;

// --- Global Variables ---
FILE* log_file = NULL;
StatsList mac_stats;
StatsList ip_stats;
time_t last_stat_time = 0;
int is_first_log = 1; // Flag to write header only once

// --- Function Prototypes ---
void packet_handler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data);
void initialize_stats(StatsList* list);
void update_stats(StatsList* list, const char* addr, unsigned long long len, int is_source);
void print_and_reset_stats(StatsList* list, const char* type);
void free_stats(StatsList* list);
void format_mac(const u_char* mac, char* buffer, size_t buffer_size);
void write_log_header();

// --- Main Function ---
int main() {
    int inum;
    pcap_if_t* alldevs;
    pcap_if_t* d;
    char errbuf[PCAP_ERRBUF_SIZE];
    pcap_t* adhandle;
    int i = 0;
    u_int netmask;
    // Filter changed to capture all IP packets for logging/stats purposes
    char packet_filter[] = "ip";
    struct bpf_program fcode;
    WSADATA wsaData;

    // Initialize Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "WSAStartup failed.\n");
        return -1;
    }

    // --- Device Discovery and Selection (same as previous code) ---
    if (pcap_findalldevs_ex((char*)PCAP_SRC_IF_STRING, NULL, &alldevs, errbuf) == -1) {
        fprintf(stderr, "Error finding devices: %s\n", errbuf);
        WSACleanup();
        exit(1);
    }
    // ... (device listing loop) ...
     for (d = alldevs; d; d = d->next) {
        printf("%d. %s", ++i, d->name);
        if (d->description)
            printf(" (%s)\n", d->description);
        else
            printf(" (No description available)\n");
    }

    if (i == 0) {
        printf("\nNo interfaces found! Make sure WinPcap is installed.\n");
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }

    printf("Enter the interface number (1-%d):", i);
    if (scanf("%d", &inum) != 1 || inum < 1 || inum > i) {
         printf("\nInterface number out of range.\n");
         pcap_freealldevs(alldevs);
         WSACleanup();
         return -1;
    }
    // Clear the input buffer
    while (getchar() != '\n');

    for (d = alldevs, i = 0; i < inum - 1; d = d->next, i++);
    // --- End Device Selection ---

    // --- Open Device ---
    if ((adhandle = pcap_open(d->name, 65536, PCAP_OPENFLAG_PROMISCUOUS, 1000, NULL, errbuf)) == NULL) {
        fprintf(stderr, "\nUnable to open the adapter. %s is not supported by WinPcap.\n", d->name);
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }

    // --- Check Data Link Type ---
    if (pcap_datalink(adhandle) != DLT_EN10MB) {
        fprintf(stderr, "\nThis program works only on Ethernet networks.\n");
        pcap_close(adhandle); // Close handle before freeing devs
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }

    // --- Compile and Set Filter ---
    if (d->addresses != NULL && d->addresses->netmask != NULL)
        // Correctly deference pointers to get the address
        netmask = ((struct sockaddr_in*)(d->addresses->netmask))->sin_addr.s_addr;
    else
        netmask = PCAP_NETMASK_UNKNOWN; // Use PCAP_NETMASK_UNKNOWN if no netmask found

    if (pcap_compile(adhandle, &fcode, packet_filter, 1, netmask) < 0) {
        fprintf(stderr, "\nUnable to compile the packet filter: %s\n", pcap_geterr(adhandle));
        pcap_close(adhandle);
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }

    if (pcap_setfilter(adhandle, &fcode) < 0) {
        fprintf(stderr, "\nError setting the filter: %s\n", pcap_geterr(adhandle));
        pcap_freealldevs(alldevs); // Free compiled code too
        pcap_freecode(&fcode);
        pcap_close(adhandle);
        WSACleanup();
        return -1;
    }
     pcap_freecode(&fcode); // Free compiled code after setting it

    // --- Initialize Log File and Stats ---
    log_file = fopen(LOG_FILENAME, "a"); // Open in append mode
    if (log_file == NULL) {
        perror("Error opening log file");
        pcap_close(adhandle);
        pcap_freealldevs(alldevs);
        WSACleanup();
        return -1;
    }
    // Check if file is empty to write header
    fseek(log_file, 0, SEEK_END);
    if (ftell(log_file) == 0) {
        is_first_log = 1; // Should already be 1, but for clarity
        write_log_header();
    } else {
        is_first_log = 0;
    }
    fflush(log_file); // Ensure header (if written) is flushed

    initialize_stats(&mac_stats);
    initialize_stats(&ip_stats);
    last_stat_time = time(NULL); // Initialize statistics timestamp

    printf("\nListening on %s...\nPress Ctrl+C to stop.\n", d->description);
    pcap_freealldevs(alldevs); // No longer need the device list

    // --- Start Capture Loop ---
    pcap_loop(adhandle, 0, packet_handler, NULL); // 0 means loop indefinitely

    // --- Cleanup (won't be reached in normal Ctrl+C exit, needs signal handling for proper cleanup) ---
    printf("\nCapture finished.\n");
    if (log_file != NULL) {
        fclose(log_file);
    }
    free_stats(&mac_stats);
    free_stats(&ip_stats);
    pcap_close(adhandle); // Close the capture handle
    WSACleanup();
    return 0;
}

// --- Packet Handler Callback ---
void packet_handler(u_char* param, const struct pcap_pkthdr* header, const u_char* pkt_data) {
    struct tm* ltime;
    char timestr[TIMESTAMP_STR_LEN];
    char src_mac_str[MAC_ADDR_STR_LEN];
    char dst_mac_str[MAC_ADDR_STR_LEN];
    char src_ip_str[IP_ADDR_STR_LEN];
    char dst_ip_str[IP_ADDR_STR_LEN];
    mac_header* mh;
    ip_header* ih;
    time_t local_tv_sec;
    unsigned int ip_header_length;

    // --- Timestamp ---
    local_tv_sec = header->ts.tv_sec;
    ltime = localtime(&local_tv_sec);
    if (ltime != NULL) {
        strftime(timestr, sizeof(timestr), "%Y-%m-%d %H:%M:%S", ltime);
    } else {
        strncpy(timestr, "Time Error", sizeof(timestr) - 1);
        timestr[sizeof(timestr)-1] = '\0';
    }


    // --- Parse Headers ---
    // Assume Ethernet II frames
    mh = (mac_header*)pkt_data;

    // Check if it's an IP packet (EtherType 0x0800)
    if (ntohs(mh->type) == 0x0800) {
        ih = (ip_header*)(pkt_data + sizeof(mac_header)); // Point to start of IP header

        // Calculate actual IP header length (IHL field * 4 bytes)
        ip_header_length = (ih->ver_ihl & 0x0F) * 4;

        // Basic validation: check minimum IP header length and if packet is long enough
        if (ip_header_length < 20 || header->caplen < (sizeof(mac_header) + ip_header_length)) {
           // Malformed or truncated IP header, skip further IP processing
           return;
        }


        // --- Extract Information ---
        format_mac(mh->src_addr, src_mac_str, sizeof(src_mac_str));
        format_mac(mh->dest_addr, dst_mac_str, sizeof(dst_mac_str));

        // Use inet_ntoa for IP addresses (simpler than manual formatting)
        strncpy(src_ip_str, inet_ntoa(ih->saddr), sizeof(src_ip_str) - 1);
        src_ip_str[sizeof(src_ip_str)-1] = '\0'; // Ensure null termination
        strncpy(dst_ip_str, inet_ntoa(ih->daddr), sizeof(dst_ip_str) - 1);
        dst_ip_str[sizeof(dst_ip_str)-1] = '\0'; // Ensure null termination

        unsigned int frame_len = header->len; // Original length of the packet

        // --- Log to CSV File ---
        if (log_file != NULL) {
            fprintf(log_file, "%s,%s,%s,%s,%s,%u\n",
                    timestr,
                    src_mac_str, src_ip_str,
                    dst_mac_str, dst_ip_str,
                    frame_len);
            fflush(log_file); // Flush buffer to ensure data is written promptly
        }

        // --- Update Statistics ---
        update_stats(&mac_stats, src_mac_str, frame_len, 1); // 1 for source
        update_stats(&mac_stats, dst_mac_str, frame_len, 0); // 0 for destination
        update_stats(&ip_stats, src_ip_str, frame_len, 1);
        update_stats(&ip_stats, dst_ip_str, frame_len, 0);

        // --- Check for Periodic Statistics Report ---
        time_t current_time = header->ts.tv_sec;
        if (current_time >= last_stat_time + STATS_INTERVAL_SECONDS) {
             printf("\n--- Statistics Report (%s) ---\n", timestr);
             print_and_reset_stats(&mac_stats, "MAC Address");
             print_and_reset_stats(&ip_stats, "IP Address");
             printf("----------------------------------------\n\n");
             last_stat_time = current_time; // Reset timer for the next interval
        }
    }
    // Add handling for other EtherTypes like ARP (0x0806) if needed
}

// --- Helper Functions ---

// Initialize a StatsList
void initialize_stats(StatsList* list) {
    list->stats = NULL;
    list->count = 0;
    list->capacity = 0;
}

// Update statistics for a given address
void update_stats(StatsList* list, const char* addr, unsigned long long len, int is_source) {
    size_t i;
    // Find existing entry
    for (i = 0; i < list->count; ++i) {
        if (strcmp(list->stats[i].address, addr) == 0) {
            if (is_source) {
                list->stats[i].sent_bytes += len;
            } else {
                list->stats[i].recv_bytes += len;
            }
            return; // Found and updated
        }
    }

    // Not found, add new entry
    // Check capacity and resize if needed
    if (list->count >= list->capacity) {
        size_t new_capacity = (list->capacity == 0) ? 10 : list->capacity * 2; // Initial capacity 10, then double
        AddressStat* new_stats = (AddressStat*)realloc(list->stats, new_capacity * sizeof(AddressStat));
        if (new_stats == NULL) {
            perror("Failed to reallocate memory for stats");
            // In a real app, might handle this more gracefully
            return;
        }
        list->stats = new_stats;
        list->capacity = new_capacity;
    }

    // Add the new stat
    strncpy(list->stats[list->count].address, addr, sizeof(list->stats[0].address) - 1);
    list->stats[list->count].address[sizeof(list->stats[0].address) - 1] = '\0'; // Ensure null termination
    if (is_source) {
        list->stats[list->count].sent_bytes = len;
        list->stats[list->count].recv_bytes = 0;
    } else {
        list->stats[list->count].sent_bytes = 0;
        list->stats[list->count].recv_bytes = len;
    }
    list->count++;
}

// Print current statistics and reset counters
void print_and_reset_stats(StatsList* list, const char* type) {
    printf("--- %s Statistics ---\n", type);
    printf("%-20s | %-18s | %-18s\n", type, "Bytes Sent", "Bytes Received");
    printf("----------------------------------------------------------\n");
    if (list->count == 0) {
        printf("  (No activity recorded in this interval)\n");
    } else {
        for (size_t i = 0; i < list->count; ++i) {
             // Only print if there was activity
             if (list->stats[i].sent_bytes > 0 || list->stats[i].recv_bytes > 0) {
                 printf("%-20s | %-18llu | %-18llu\n",
                        list->stats[i].address,
                        list->stats[i].sent_bytes,
                        list->stats[i].recv_bytes);
            }
            // Reset counters for the next interval
            list->stats[i].sent_bytes = 0;
            list->stats[i].recv_bytes = 0;
        }
    }
     printf("----------------------------------------------------------\n");
     // Optional: Shrink the list if count is much smaller than capacity (more complex)
     // Optional: Remove entries with zero activity (more complex)
}

// Free memory allocated for a StatsList
void free_stats(StatsList* list) {
    if (list->stats != NULL) {
        free(list->stats);
        list->stats = NULL;
        list->count = 0;
        list->capacity = 0;
    }
}

// Format MAC address into XX-XX-XX-XX-XX-XX string
void format_mac(const u_char* mac, char* buffer, size_t buffer_size) {
    snprintf(buffer, buffer_size, "%02X-%02X-%02X-%02X-%02X-%02X",
             mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
}

// Write the CSV header to the log file
void write_log_header() {
    if (log_file != NULL && is_first_log) {
        fprintf(log_file, "Timestamp,Source MAC,Source IP,Destination MAC,Destination IP,Frame Length\n");
        fflush(log_file); // Ensure header is written
        is_first_log = 0; // Clear the flag
    }
}